我們昨天從零開始手寫了 LLaMA2，並且體驗了 tokenizer 訓練和簡單的預訓練流程，雖然收穫滿滿，但實際上很少人會真的從零實作 Transformer 再來訓練，因為每次都要自己寫 Attention、LayerNorm 等模塊，太花時間，也很難針對多張顯卡高效率的訓練，另外也無法快速的載入 Hugging Face 已有的開源 LLM 權重與社群資源。
　　因此，今天我們要使用 Hugging Face Transformers + DeepSpeed 來完成 LLM 的預訓練！

1.為什麼要用 Transformers？

　　Hugging Face 的 Transformers 框架，已經封裝了 BERT、GPT、LLaMA、T5 等各種主流架構，不管在模型載入、分散式訓練、資料集等，都可以很方便的使用，我們善用框架和開源資源，把訓練和應用做的更快更好。

2.初始化模型與 Tokenizer

假設我們要使用 Qwen-2.5-1.5B，可以直接載入 config 與 tokenizer

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

# 載入模型設定
config = AutoConfig.from_pretrained("Qwen/Qwen2.5-1.5B")

# 初始化模型（從零開始）
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)

# 或者直接載入預訓練權重
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B", trust_remote_code=True)

# 載入對應的 tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B")

3.準備預訓練數據

　　和昨天一樣，我們使用 Mobvoi Seq-Monkey 資料集作為例子，在 Hugging Face 中，可以透過 datasets 直接載入 JSON 格式：

from datasets import load_dataset

ds = load_dataset('json', data_files='mobvoi_seq_monkey_general_open_corpus.jsonl')
print(ds["train"][0])  # 看第一筆

接下來用 tokenizer 轉換成模型能吃的格式：

def tokenize_function(examples):
    return tokenizer(examples["text"])

tokenized_datasets = ds.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"],
    num_proc=10,
    desc="Tokenizing dataset",
)

預訓練通常需要長序列（2048 tokens），我們把多段文本拼接起來：

from itertools import chain

block_size = 2048

def group_texts(examples):
    concatenated = {k: list(chain(*examples[k])) for k in examples.keys()}
    total_length = len(concatenated["input_ids"])
    total_length = (total_length // block_size) * block_size
    result = {
        k: [t[i:i+block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result

lm_datasets = tokenized_datasets.map(
    group_texts, batched=True, num_proc=10, desc="Grouping texts"
)
train_dataset = lm_datasets["train"]

4.用 Trainer 開始訓練

Transformers 有封裝好一個 Trainer 類別，能直接讓我們開始訓練：

from transformers import Trainer, TrainingArguments, default_data_collator

training_args = TrainingArguments(
    output_dir="output",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    logging_steps=10,
    save_steps=100,
    learning_rate=1e-4,
    gradient_checkpointing=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
    data_collator=default_data_collator,
)

trainer.train()

要注意！這個訓練需要花費大量硬體資源

參考連結：
https://datawhalechina.github.io/happy-llm/#/